import pygame
import numpy as np
import random
import math
import argparse
from PIL import Image
import sys
from pygame import gfxdraw

class Particle:
    def __init__(self, x, y, target_x=None, target_y=None):
        self.x = x
        self.y = y
        self.vx = random.uniform(-2, 2)
        self.vy = random.uniform(-2, 2)
        self.target_x = target_x
        self.target_y = target_y
        self.color = (255, 255, 255)
        self.size = 2
        self.transition_progress = 0.0
        
    def update_boids(self, particles, width, height):
        separation = self.separation(particles)
        alignment = self.alignment(particles)
        cohesion = self.cohesion(particles)
        
        self.vx += separation[0] * 1.5 + alignment[0] * 1.0 + cohesion[0] * 1.0
        self.vy += separation[1] * 1.5 + alignment[1] * 1.0 + cohesion[1] * 1.0
        
        speed = math.sqrt(self.vx**2 + self.vy**2)
        if speed > 3:
            self.vx = (self.vx / speed) * 3
            self.vy = (self.vy / speed) * 3
        
        self.x += self.vx
        self.y += self.vy
        
        if self.x < 0 or self.x > width:
            self.vx *= -1
        if self.y < 0 or self.y > height:
            self.vy *= -1
            
        self.x = max(0, min(width, self.x))
        self.y = max(0, min(height, self.y))
    
    def update_boids_optimized(self, nearby_particles, width, height):
        """优化的boids更新，使用预计算的邻近粒子"""
        if not nearby_particles:
            # 如果没有邻近粒子，随机游走
            self.vx += random.uniform(-0.1, 0.1)
            self.vy += random.uniform(-0.1, 0.1)
        else:
            separation = self.separation(nearby_particles)
            alignment = self.alignment(nearby_particles)
            cohesion = self.cohesion(nearby_particles)
            
            self.vx += separation[0] * 1.2 + alignment[0] * 0.8 + cohesion[0] * 0.8
            self.vy += separation[1] * 1.2 + alignment[1] * 0.8 + cohesion[1] * 0.8
        
        # 速度限制
        speed = math.sqrt(self.vx**2 + self.vy**2)
        max_speed = 2.5
        if speed > max_speed:
            self.vx = (self.vx / speed) * max_speed
            self.vy = (self.vy / speed) * max_speed
        
        self.x += self.vx
        self.y += self.vy
        
        # 边界处理
        margin = 20
        if self.x < margin:
            self.vx += 0.1
        elif self.x > width - margin:
            self.vx -= 0.1
        if self.y < margin:
            self.vy += 0.1
        elif self.y > height - margin:
            self.vy -= 0.1
            
        self.x = max(0, min(width, self.x))
        self.y = max(0, min(height, self.y))
    
    def separation(self, particles):
        steer_x, steer_y = 0, 0
        count = 0
        
        for other in particles:
            if other != self:
                distance = math.sqrt((self.x - other.x)**2 + (self.y - other.y)**2)
                if distance < 25 and distance > 0:
                    diff_x = self.x - other.x
                    diff_y = self.y - other.y
                    diff_x /= distance
                    diff_y /= distance
                    steer_x += diff_x
                    steer_y += diff_y
                    count += 1
        
        if count > 0:
            steer_x /= count
            steer_y /= count
        
        return (steer_x, steer_y)
    
    def alignment(self, particles):
        avg_vx, avg_vy = 0, 0
        count = 0
        
        for other in particles:
            if other != self:
                distance = math.sqrt((self.x - other.x)**2 + (self.y - other.y)**2)
                if distance < 50:
                    avg_vx += other.vx
                    avg_vy += other.vy
                    count += 1
        
        if count > 0:
            avg_vx /= count
            avg_vy /= count
            return (avg_vx - self.vx, avg_vy - self.vy)
        
        return (0, 0)
    
    def cohesion(self, particles):
        center_x, center_y = 0, 0
        count = 0
        
        for other in particles:
            if other != self:
                distance = math.sqrt((self.x - other.x)**2 + (self.y - other.y)**2)
                if distance < 50:
                    center_x += other.x
                    center_y += other.y
                    count += 1
        
        if count > 0:
            center_x /= count
            center_y /= count
            return ((center_x - self.x) * 0.01, (center_y - self.y) * 0.01)
        
        return (0, 0)
    
    def update_transition(self, target_x, target_y, speed=0.02):
        if self.target_x is not None and self.target_y is not None:
            dx = target_x - self.x
            dy = target_y - self.y
            
            self.x += dx * speed
            self.y += dy * speed
            
            self.vx *= 0.95
            self.vy *= 0.95
            
            distance = math.sqrt(dx**2 + dy**2)
            if distance < 5:
                self.x = target_x
                self.y = target_y
    
    def draw(self, screen):
        # 确保颜色值在有效范围内
        color = tuple(max(0, min(255, c)) for c in self.color)
        # 使用普通圆形绘制避免颜色错误
        pygame.draw.circle(screen, color, (int(self.x), int(self.y)), self.size)

class ParticleAnimation:
    def __init__(self, image_path=None, width=800, height=600, num_particles=500):
        pygame.init()
        self.width = width
        self.height = height
        self.screen = pygame.display.set_mode((width, height))
        pygame.display.set_caption("粒子动画 - 鸟群模拟")
        
        self.clock = pygame.time.Clock()
        self.particles = []
        self.num_particles = min(num_particles, 800)  # 限制最大粒子数
        self.mode = "boids"  # "boids" or "transition"
        self.transition_start_time = None
        self.boids_duration = 3000  # 3秒鸟群模式
        self.has_image = image_path is not None
        
        # 性能优化
        self.spatial_grid = {}
        self.grid_size = 50
        self.fps_counter = 0
        self.last_fps_time = pygame.time.get_ticks()
        self.current_fps = 0
        
        if self.has_image:
            self.load_image_particles(image_path)
            self.mode = "boids"  # 有图片时先显示鸟群，然后过渡到照片
        else:
            self.mode = "boids_only"  # 无图片时只显示鸟群
            
        self.init_boids_particles()
        
    def load_image_particles(self, image_path):
        try:
            image = Image.open(image_path)
            # 降低分辨率以提高性能
            image = image.resize((self.width//2, self.height//2))
            image = image.convert('RGB')
            
            self.target_particles = []
            pixels = np.array(image)
            
            # 更智能的采样
            step = max(2, int(math.sqrt((self.width * self.height) / self.num_particles)))
            
            # 采样重要像素
            for i in range(0, self.height//2, step):
                for j in range(0, self.width//2, step):
                    if len(self.target_particles) < self.num_particles:
                        r, g, b = pixels[i, j]
                        brightness = (r + g + b) / 3
                        if brightness > 50:  # 提高亮度阈值
                            self.target_particles.append({
                                'x': j * 2,  # 放大回原尺寸
                                'y': i * 2,
                                'color': (r, g, b)
                            })
            
            # 如果粒子数量不够，随机补充
            while len(self.target_particles) < self.num_particles:
                x = random.randint(0, self.width)
                y = random.randint(0, self.height)
                color = (random.randint(100, 255), random.randint(100, 255), random.randint(100, 255))
                self.target_particles.append({'x': x, 'y': y, 'color': color})
                
        except Exception as e:
            print(f"无法加载图片: {e}")
            self.create_default_target()
    
    def create_default_target(self):
        self.target_particles = []
        for i in range(self.num_particles):
            x = random.randint(0, self.width)
            y = random.randint(0, self.height)
            color = (random.randint(100, 255), random.randint(100, 255), random.randint(100, 255))
            self.target_particles.append({'x': x, 'y': y, 'color': color})
    
    def init_boids_particles(self):
        for i in range(self.num_particles):
            x = random.randint(0, self.width)
            y = random.randint(0, self.height)
            particle = Particle(x, y)
            self.particles.append(particle)
    
    def update_spatial_grid(self):
        """更新空间网格以优化邻近粒子查找"""
        self.spatial_grid.clear()
        for particle in self.particles:
            grid_x = int(particle.x // self.grid_size)
            grid_y = int(particle.y // self.grid_size)
            key = (grid_x, grid_y)
            if key not in self.spatial_grid:
                self.spatial_grid[key] = []
            self.spatial_grid[key].append(particle)
    
    def get_nearby_particles(self, particle, radius):
        """获取邻近粒子"""
        grid_x = int(particle.x // self.grid_size)
        grid_y = int(particle.y // self.grid_size)
        
        nearby = []
        # 检查周围9个网格
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                key = (grid_x + dx, grid_y + dy)
                if key in self.spatial_grid:
                    for other in self.spatial_grid[key]:
                        if other != particle:
                            distance = math.sqrt((particle.x - other.x)**2 + (particle.y - other.y)**2)
                            if distance <= radius:
                                nearby.append(other)
        return nearby
    
    def update(self):
        current_time = pygame.time.get_ticks()
        
        # 更新FPS计数器
        self.fps_counter += 1
        if current_time - self.last_fps_time >= 1000:
            self.current_fps = self.fps_counter
            self.fps_counter = 0
            self.last_fps_time = current_time
        
        if self.mode == "boids":
            # 更新空间网格
            self.update_spatial_grid()
            
            for particle in self.particles:
                # 使用优化的邻近粒子查找
                nearby_particles = self.get_nearby_particles(particle, 50)
                particle.update_boids_optimized(nearby_particles, self.width, self.height)
            
            if current_time > self.boids_duration and self.has_image:
                self.mode = "transition"
                self.transition_start_time = current_time
                
        elif self.mode == "boids_only":
            # 纯鸟群模式，持续更新
            self.update_spatial_grid()
            
            for particle in self.particles:
                nearby_particles = self.get_nearby_particles(particle, 50)
                particle.update_boids_optimized(nearby_particles, self.width, self.height)
                
        elif self.mode == "transition":
            progress = min(1.0, (current_time - self.transition_start_time) / 2000.0)
            
            for i, particle in enumerate(self.particles):
                if i < len(self.target_particles):
                    target = self.target_particles[i]
                    
                    particle.update_transition(target['x'], target['y'], 0.08)
                    
                    # 颜色过渡
                    start_color = (255, 255, 255)
                    end_color = target['color']
                    
                    r = int(start_color[0] + (end_color[0] - start_color[0]) * progress)
                    g = int(start_color[1] + (end_color[1] - start_color[1]) * progress)
                    b = int(start_color[2] + (end_color[2] - start_color[2]) * progress)
                    
                    particle.color = (r, g, b)
    
    def draw(self):
        self.screen.fill((0, 0, 0))
        
        # 批量绘制粒子以提高性能
        for particle in self.particles:
            particle.draw(self.screen)
        
        # 显示性能信息
        font = pygame.font.Font(None, 24)
        
        if self.mode == "boids_only":
            mode_text = "鸟群模拟模式"
        elif self.mode == "boids":
            mode_text = "鸟群模式"
        elif self.mode == "transition":
            mode_text = "照片形成中..."
        else:
            mode_text = "未知模式"
            
        fps_text = f"FPS: {self.current_fps}"
        particles_text = f"粒子数: {len(self.particles)}"
        
        text1 = font.render(mode_text, True, (255, 255, 255))
        text2 = font.render(fps_text, True, (255, 255, 255))
        text3 = font.render(particles_text, True, (255, 255, 255))
        
        self.screen.blit(text1, (10, 10))
        self.screen.blit(text2, (10, 35))
        self.screen.blit(text3, (10, 60))
        
        pygame.display.flip()
    
    def run(self):
        running = True
        
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_r:
                        self.__init__(self.target_particles if hasattr(self, 'target_particles') else None, 
                                     self.width, self.height, self.num_particles)
                    elif event.key == pygame.K_SPACE:
                        self.mode = "transition"
                        self.transition_start_time = pygame.time.get_ticks()
            
            self.update()
            self.draw()
            self.clock.tick(60)
        
        pygame.quit()
        sys.exit()

def main():
    parser = argparse.ArgumentParser(description='粒子动画程序 - 鸟群模拟')
    parser.add_argument('--image', help='照片文件路径（可选）')
    parser.add_argument('--width', type=int, default=800, help='窗口宽度')
    parser.add_argument('--height', type=int, default=600, help='窗口高度')
    parser.add_argument('--particles', type=int, default=500, help='粒子数量')
    
    args = parser.parse_args()
    
    animation = ParticleAnimation(args.image, args.width, args.height, args.particles)
    animation.run()

if __name__ == "__main__":
    main()